package util

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
)

// 生成密钥对
func GenerateRsaKey(keySize int) {
	//创建私钥文件
	//使用rsa中的GenerateKey方法生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		panic(err)
	}
	//通过x509标准将得到的ras私钥序列化为ASN.1 的 DER编码字符串
	key509pri := x509.MarshalPKCS1PrivateKey(privateKey)
	//将私钥字符串设置到pem格式块中初始化一个pem.Block块
	block := pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: key509pri,
	}
	file, err := os.Create("private.pem")
	if err != nil {
		panic(err)
	}
	pem.Encode(file, &block)
	file.Close()
	fmt.Println("private.pem is Create")
	//创建公钥文件
	//从得到的私钥对象中将公钥信息取出
	publicKey := privateKey.PublicKey
	//通过x509标准将得到 的rsa公钥序列化为字符串
	key509pub, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		panic(err)
	}
	//将公钥字符串设置到pem格式块中
	block = pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: key509pub,
	}
	//通过pem将设置好的数据进行编码, 并写入磁盘文件
	file, err = os.Create("public.pem")
	if err != nil {
		panic(err)
	}
	pem.Encode(file, &block)
	file.Close()
	fmt.Println("public.pem is Create")
}

// rsa 公钥从文件中提取
func RSAFromFileGetPublicKey(fileName string) *rsa.PublicKey {
	//将公钥文件中的公钥读出, 得到使用pem编码的字符串-- 读文件
	file, err := os.Open(fileName)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	fileInfo, err := file.Stat()
	if err != nil {
		panic(err)
	}
	bytes := make([]byte, fileInfo.Size())
	_, err = file.Read(bytes)
	if err != nil {
		panic(err)
	}
	//将得到的字符串解码-- pem.Decode
	decodePubKey, _ := pem.Decode(bytes)
	//使用x509将编码之后的公钥解析出来
	publicKey, err := x509.ParsePKIXPublicKey(decodePubKey.Bytes)
	if err != nil {
		panic(err)
	}
	//断言转换类型
	key := publicKey.(*rsa.PublicKey)
	return key
}

// rsa 公钥加密
func RSAEncrypt(plainText []byte, fileName string) []byte {
	key := RSAFromFileGetPublicKey(fileName)
	//使用得到的公钥通过rsa进行数据加密
	cipherText, err := rsa.EncryptPKCS1v15(rand.Reader, key, plainText)
	if err != nil {
		panic(err)
	}
	return cipherText
}

// rsa 私钥从文件中提取
func RSAFromFileGetPrivateKey(fileName string) *rsa.PrivateKey {
	//将私钥文件中的私钥读出, 得到使用pem编码的字符串-- 读文件
	file, err := os.Open(fileName)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	fileInfo, err := file.Stat()
	if err != nil {
		panic(err)
	}
	bytes := make([]byte, fileInfo.Size())
	_, err = file.Read(bytes)
	if err != nil {
		panic(err)
	}
	//将得到的字符串解码-- pem.Decode
	decodePriKey, _ := pem.Decode(bytes)
	//使用x509将编码之后的公钥解析出来
	privateKey, err := x509.ParsePKCS1PrivateKey(decodePriKey.Bytes)
	if err != nil {
		panic(err)
	}
	return privateKey
}

// rsa 私钥解密
func RSADecrypt(cipherText []byte, fileName string) []byte {
	//文件提取私钥
	privateKey := RSAFromFileGetPrivateKey(fileName)
	//使用得到的公钥通过rsa进行数据加密
	plainText, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, cipherText)
	if err != nil {
		panic(err)
	}
	return plainText
}
